#include <am.h>
#include <nemu.h>
#include <klib.h>

static AddrSpace kas = {};
static void* (*pgalloc_usr)(int) = NULL;
static void (*pgfree_usr)(void*) = NULL;
static int vme_enable = 0;

static Area segments[] = {      // Kernel memory mappings
  NEMU_PADDR_SPACE
};

#define USER_SPACE RANGE(0x40000000, 0x80000000)

static inline void set_satp(void *pdir) {
  uintptr_t mode = 1ul << (__riscv_xlen - 1);
  asm volatile("csrw satp, %0" : : "r"(mode | ((uintptr_t)pdir >> 12)));
}

static inline uintptr_t get_satp() {
  uintptr_t satp;
  asm volatile("csrr %0, satp" : "=r"(satp));
  return satp << 12;
}

bool vme_init(void* (*pgalloc_f)(int), void (*pgfree_f)(void*)) {
  pgalloc_usr = pgalloc_f;
  pgfree_usr = pgfree_f;

  kas.ptr = pgalloc_f(PGSIZE);

  int i;
  for (i = 0; i < LENGTH(segments); i ++) {
    void *va = segments[i].start;
    for (; va < segments[i].end; va += PGSIZE) {
      map(&kas, va, va, 0);
    }
  }

  set_satp(kas.ptr);
  vme_enable = 1;

  return true;
}

void protect(AddrSpace *as) {
  PTE *updir = (PTE*)(pgalloc_usr(PGSIZE));//printf("%08x\n",updir);
  as->ptr = updir;
  as->area = USER_SPACE;
  as->pgsize = PGSIZE;
  // map kernel space
  memcpy(updir, kas.ptr, PGSIZE);
}

void unprotect(AddrSpace *as) {
}

void __am_get_cur_as(Context *c) {//printf("%08x\n",c);
  c->pdir = (vme_enable ? (void *)get_satp() : NULL);
}

void __am_switch(Context *c) {
  if (vme_enable&&c->pdir!=NULL) {
    set_satp(c->pdir);
  }
}

uintptr_t check_map(AddrSpace *as,void * va){
  uintptr_t high=(uintptr_t)va>>22,low=(uintptr_t)va>>12&0x3ff;
  uintptr_t * pos=(uintptr_t *)as->ptr+high;
  if(*pos==0) return 0;
  uintptr_t * pos2=(uintptr_t *)(*pos&~0x3ffu);
  return(pos2[low]);
}

void map(AddrSpace *as, void *va, void *pa, int prot) {
  uintptr_t high=(uintptr_t)va>>22,low=(uintptr_t)va>>12&0x3ff;
  uintptr_t * pos=(uintptr_t *)as->ptr+high;
  if(*pos==0) *pos=(uintptr_t)pgalloc_usr(PGSIZE);
  uintptr_t * pos2=(uintptr_t *)(*pos&~0x3ffu);
  assert(pos2[low]==0);
  pos2[low]=(uintptr_t)pa;
//  if((uintptr_t)va<0x80000000u) printf("%08x->%08x\n",va,pa);
  return;
}

Context *ucontext(AddrSpace *as, Area kstack, void *entry) {
  assert(entry!=NULL);
  Context * ret=(Context *)((char *)kstack.end-sizeof(Context));
  memset(ret,0,sizeof(ret));
  ret->mepc=(uintptr_t)entry;
  ret->mstatus=0x1880;
  ret->mcause=11;
  ret->pdir=as->ptr;
  ret->np=1;
  return ret;
}
